#include "trace_custom_tiling.h"
#include "data_utils.h"
#include <sys/stat.h>
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_trace_custom.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void trace_custom(GM_ADDR x, GM_ADDR y, TraceCustomTilingData tiling);
#endif
extern void GenerateTilingData(uint8_t *tilingBuf, uint32_t matrixNum);

static size_t GetFileSize(const char *path)
{
    struct stat st{};
    if (stat(path, &st) != 0) return 0;
    return (size_t)st.st_size;
}

int32_t main(int32_t argc, char *argv[])
{
    const char *inputPath = "./input/input_x.bin";
    size_t singleMatrixBytes = 777ULL * 777ULL * sizeof(uint16_t);
    // 先生成数据脚本后才有文件
    // run.sh 已先调用 gen_data.py
    size_t fileSize = GetFileSize(inputPath);
    if (fileSize == 0 || fileSize % singleMatrixBytes != 0) {
        printf("[ERROR] input file size invalid: %zu\n", fileSize);
        return -1;
    }
    uint32_t matrixNum = static_cast<uint32_t>(fileSize / singleMatrixBytes);

    uint32_t blockDim = matrixNum;
    size_t tilingSize = sizeof(TraceCustomTilingData);
    size_t inputByteSize  = fileSize;                    // 全部矩阵
    size_t outputByteSize = matrixNum * sizeof(uint16_t);

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *tiling = (uint8_t *)AscendC::GmAlloc(tilingSize);
    GenerateTilingData(tiling, matrixNum);
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *y = (uint8_t *)AscendC::GmAlloc(outputByteSize);
    ReadFile(inputPath, inputByteSize, x, inputByteSize);
    for (size_t i = 0; i < outputByteSize; ++i) y[i] = 0;
    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(trace_custom, blockDim, x, y, *reinterpret_cast<TraceCustomTilingData *>(tiling));
    WriteFile("./output/output_y.bin", y, outputByteSize);
    AscendC::GmFree(x); AscendC::GmFree(y); AscendC::GmFree(tiling);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *tiling = nullptr;
    uint8_t *xHost=nullptr,*yHost=nullptr;
    uint8_t *xDevice=nullptr,*yDevice=nullptr;

    CHECK_ACL(aclrtMallocHost((void **)&tiling, tilingSize));
    GenerateTilingData(tiling, matrixNum);

    CHECK_ACL(aclrtMallocHost((void **)&xHost, inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)&yHost, outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&yDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMemset(yDevice, outputByteSize, 0, outputByteSize));

    ReadFile(inputPath, inputByteSize, xHost, inputByteSize);
    CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    ACLRT_LAUNCH_KERNEL(trace_custom)(blockDim, stream, xDevice, yDevice,
        reinterpret_cast<TraceCustomTilingData *>(tiling));
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(yHost, outputByteSize, yDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_y.bin", yHost, outputByteSize);

    CHECK_ACL(aclrtFree(xDevice)); CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFreeHost(xHost)); CHECK_ACL(aclrtFreeHost(yHost));
    CHECK_ACL(aclrtFreeHost(tiling));
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}